package com.demo.agent.firstagent;

import javassist.*;
import javassist.expr.ExprEditor;
import javassist.expr.MethodCall;

import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.security.ProtectionDomain;

/**
 * @author 许路路
 */
public class FirstAgent implements ClassFileTransformer {

    public final String injectedClassName = "com.demo.agent.test.DemoClass";
    public final String methodName = "hello";

    @Override
    public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {
        className = className.replace("/", ".");
        //System.out.println(++times + ":类：" + className);
        if (className.equals(injectedClassName)) {
            CtClass ctclass;
            try {
                // 使用全称,用于取得字节码类<使用javassist>
                ctclass = ClassPool.getDefault().get(className);
                // 得到这方法实例
                CtBehavior[] methods = ctclass.getDeclaredBehaviors();
                for (CtBehavior method : methods) {
                    enhanceMethod(method);
                }

                CtMethod ctmethod = ctclass.getDeclaredMethod(methodName);
                ctmethod.insertBefore("long start = System.nanoTime();\n System.out.println(\"开始：\" + start);");
                ctmethod.insertAfter("long end = System.nanoTime();\n System.out.println(\"结束：\"+ end);");
                return ctclass.toBytecode();
            } catch (Exception e) {
                System.out.println(e.getMessage());
                e.printStackTrace();
            }
        }
        return null;
    }

    /**
     * 方法增强 统计方法调用时间
     *
     * @param method
     * @throws Exception
     */
    private void enhanceMethod(CtBehavior method) throws Exception {
        if (method.isEmpty()) {
            return;
        }
        String methodName = method.getName();
        // 不提升main方法
        if (methodName.equalsIgnoreCase("main")) {
            return;
        }

        final StringBuilder source = new StringBuilder();
        source.append("{").append("long start = System.nanoTime();\n")
                .append("$_ = $proceed($$);\n")
                .append("System.out.print(\"method:[" + methodName + "]\");").append("\n")
                .append("System.out.println(\" cost:[\" +(System.nanoTime() -start)+ \"ns]\");")
                .append("}");
        //替换原有方法
        ExprEditor editor = new ExprEditor() {
            @Override
            public void edit(MethodCall methodCall) throws CannotCompileException {
                methodCall.replace(source.toString());
            }
        };
        method.instrument(editor);
    }
}
